package GreedyAlgorithm;
use strict;
use Bio::Phylo::BioPrio::GreedySolution;
use Bio::Phylo::BioPrio::SpeciesConservationMeasures;

=head1 NAME

GreedyAlgorithm - Implements greedy algorithms for the NAP

=head1 DESCRIPTION

Implements the greedy algorithms for the base NAP as outlined 
in Steel 2005, Pardi 2005, Hartmann & Steel 2006,2007. These 
algorithms are implemented with simplicity and readability 
being the foremost concern. They also consider all possible 
solutions to a problem (which may be many for simple instances
of the NAP).

They are nowhere near as computationally efficient as possible.
For better speed a simpler tree data structure aimed at these
algorithms would be beneficial as would implementation in C.
Dropping support for multiple solutions would also deliver
significant performance gains.

However, the aim here is for clear and readable code. Anyone 
wishing to implement a fast version of these algorithms is more
than welcome to do so! 

=cut

sub solve {
    #The tree
    my $tree = shift;
    #The survival probability
    my $sp = shift;

    #Check that this relationship is actually linear
    unless ( $sp->is_linear ) {
    	return 0;
    }

#              is_constant_cost, is_survival_ratio_constant
#CASES:  a->b, cost constant, (1-b)/(1-a) = constant
#        is_zero_one_prob, tree check
#        0->1, cost variable, ultrametric

    #Greedy works
    if ($sp->is_constant_cost) {
    	
    }

}


sub solve_constant_cost {
    #The tree
    my $tree = shift;
    #The survival probability
    my $sp = shift;

	#Get the linear survival relationships 
	$sp->get_linear_survival_tree($tree);

	#The solution class
	my $solution = GreedySolution->new;

	my @species = $sp->get_names;
	
	my @priority;
	

	
	#Make a leaf hash so we can easily access the leaf nodes by name 
	my %leafhash;
	my @leaves = @{$tree->get_terminals};
	foreach (@leaves) { $leafhash{$_->get_name} = $_; }
	
	my $setsize = 0;
	
	#Perform the greedy algorithm
	while ($setsize < scalar @species) {

		my @currentsets = $solution->get_sets_n($setsize);
		
		foreach my $set (@currentsets) {
			#Set conservation status
			_set_conserve($tree,@{$set});
			#Set edge extinction probabilities
			_set_cc_edgeextinction( $tree->get_root );
			#Set PD gain for each species	
			_set_cc_pdgain( $tree );
			
			#The maximum benefit attainable
			my $max_benefit = 0;
			
			#The greedy choices and their PD gains
			my @greedychoices;
			my @pdgain;
			my @cost;
			my $max = -1;
			#Check each species to see if it is a greedy choice
			foreach(@leaves) {
				#If a new maximum gain is found
				if ($_->get_generic('gain')>$max) {
					#Set the maximum value
					$max = $_->get_generic('gain');
					#Reset the greedy choices
					@greedychoices = ($_->get_name);
					@pdgain = ($max);
					@cost = ($_->get_generic('c'));
			    #If an equally good gain is found
				} elsif ($_->get_generic('gain')==$max) {
					#Add to the greedy choices
					push(@greedychoices, $_->get_name);
					push(@pdgain, $max);
					push(@cost, $_->get_generic('c'));
				}
			}
			$solution->add_transition($set,\@greedychoices,\@pdgain, \@cost);				
		} 
		$setsize++;
	}

	return $solution;
}

#Set conservation status for all species
sub _set_conserve {
	my $tree = shift;
	my %toconserve;
	
	#Make a hash so it is easy to check if a species will be conserved
	while (@_) { $toconserve{shift(@_)} = 1; }
	
	#Get all leaf nodes	
	my @leaves = @{$tree->get_terminals};
	
	#For each leaf check wether to conserve it
	foreach (@leaves) {
		#The species will be conserved
		if ($toconserve{$_->get_name}) {
			$_->set_generic( conserved => 1 );
		#Won't be conserved 
		} else {
			$_->set_generic( conserved => 0 );
		}
	}	
}

#Sets the PD increase obtained by conserving each species
sub _set_cc_pdgain {
	my $tree = shift;
	
	my @leaves = @{$tree->get_terminals};
	
	#For each leaf store the gain
	foreach my $leaf (@leaves) {
		my $gain = 0;
		
		#If this species has not already been conserved we calculate the PD gain it provides
		unless ($leaf->get_generic('conserved')) {
			my $a = $leaf->get_generic('a');
			my $b = $leaf->get_generic('b');
			
			#Get the gain from each ancestral edge
			foreach my $ancestor (($leaf, @{$leaf->get_ancestors})) {
				$gain += $ancestor->get_branch_length * $ancestor->get_generic('extinction') * ($b-$a)/(1-$a); 
			}
		}
		
		#Store the total gain
		$leaf->set_generic(gain=>$gain);
	} 
	
}

#Sets each node's extinction probability 
sub _set_cc_edgeextinction {
	my $node = shift;
	
	#If this is a terminal node we set it's extinction probability
	if ($node->is_terminal) {
		my $name = $node->get_name;
		if ($node->get_generic('conserved')) {
			$node->set_generic( extinction => 1-$node->get_generic('b') );
		} else {
			$node->set_generic( extinction => 1-$node->get_generic('a') );
		}
		return;
	}
	
	#For non-terminals we call this routine for the children and multiply
	#their extinction probabilities together to get this node's probability
	my $extinction = 1;
	foreach my $child (@{$node->get_children}) {
		_set_cc_edgeextinction($child);
		$extinction *= $child->get_generic('extinction');
	}
	
	#Set this node's extinction probability
	$node->set_generic( extinction => $extinction);
}
1;